# -*- coding: utf-8 -*-
"""
文件名: image_enhance_cls.py

功能:
  - 对分类数据集进行数据增强，支持“正常/异常”样本的不同增强策略。
  - 扩展增广为 D4 群的 8 类唯一几何变换（去除冗余组合）：
        h(左右翻) / v(上下翻) / vh(180°) / rot90 / rot270 / d(主对角线翻转) / ad(副对角线翻转)
    * 备注：vh 等价于水平+垂直或 180° 旋转；这 7 类（不含 identity）覆盖常见“先翻再转”的唯一结果。
  - 输出：
        images/                      增强后的图片
        enhanced_annotations.json    逐条增强样本条目
        meta.json                    统计与配置

示例:
  >>> from image_enhance_cls import build_enhanced_cls
  >>> out = build_enhanced_cls(
  ...     annotations_path="./experiments/dataset/cls_annotations.json",
  ...     enhanced_save_dir="./experiments/dataset_enhanced/cls",
  ...     normal_key="正常眼底",
  ...     count_strategy="random",
  ...     random_k_choices=(1,2,3),
  ...     per_image_k=2,
  ...     # 概率字典对齐到新的操作集合，未提供的操作默认权重0:
  ...     op_probs={"h":0.25,"v":0.25,"vh":0.1,"rot90":0.2,"rot270":0.1,"d":0.05,"ad":0.05},
  ...     normal_no_aug_prob=0.7,        # 70% 正常图不增强
  ...     abnormal_no_aug_prob=0.0,      # 异常图必增强
  ...     allowed_ops=None,              # None 表示使用所有上面列出的操作
  ...     seed=42,
  ...     force_rebuild=False
  ... )
  >>> print(out["stats"])

依赖:
  - Pillow
"""

import os
import json
import shutil
import random
import hashlib
from typing import Dict, List, Tuple, Optional, Iterable

from PIL import Image


# ==============================
# 内部工具
# ==============================

def _safe_makedirs(path: str) -> None:
    os.makedirs(path, exist_ok=True)

def _safe_load_json(path: str) -> dict:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        raise RuntimeError(f"[读取JSON失败] path={path}, error={repr(e)}")

def _safe_dump_json(data: dict, path: str) -> None:
    try:
        _safe_makedirs(os.path.dirname(path))
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    except Exception as e:
        raise RuntimeError(f"[写入JSON失败] path={path}, error={repr(e)}")

def _build_new_name(image_name: str, op: str) -> str:
    """
    命名规则:
      abc.png -> abc_h.png / abc_v.png / abc_vh.png / abc_r90.png / abc_r270.png / abc_d.png / abc_ad.png
    """
    stem, ext = os.path.splitext(image_name)
    tag_map = {
        "h": "_h",
        "v": "_v",
        "vh": "_vh",
        "rot90": "_r90",
        "rot270": "_r270",
        "d": "_d",     # 主对角线翻转
        "ad": "_ad"    # 副对角线翻转
    }
    tag = tag_map.get(op, "_aug")
    return f"{stem}{tag}{ext}"

def _md5_of_file(path: str) -> str:
    h = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()

# ==============================
# 变换（Pillow）
# ==============================

# D4 群的 7 个非恒等变换（不输出 identity，identity 代表“不增强”）
# 说明：Pillow transpose 常量
#   Image.FLIP_LEFT_RIGHT -> 左右翻
#   Image.FLIP_TOP_BOTTOM -> 上下翻
#   Image.ROTATE_90      -> 逆时针90（这里我们采用顺/逆约定见下）
#   Image.ROTATE_180
#   Image.ROTATE_270
#   Image.TRANSPOSE      -> 主对角线翻转（等价于转置）
#   Image.TRANSVERSE     -> 副对角线翻转
def _apply_op(img: Image.Image, op: str) -> Image.Image:
    if op == "h":
        return img.transpose(Image.FLIP_LEFT_RIGHT)
    elif op == "v":
        return img.transpose(Image.FLIP_TOP_BOTTOM)
    elif op == "vh":
        return img.transpose(Image.ROTATE_180)  # 等价于水平再垂直
    elif op == "rot90":
        # 约定：顺时针90° = ROTATE_270（Pillow的ROTATE_270就是顺时针90）
        return img.transpose(Image.ROTATE_270)
    elif op == "rot270":
        # 顺时针270° = 逆时针90° = ROTATE_90
        return img.transpose(Image.ROTATE_90)
    elif op == "d":
        # 主对角线翻转
        return img.transpose(Image.TRANSPOSE)
    elif op == "ad":
        # 副对角线翻转
        return img.transpose(Image.TRANSVERSE)
    else:
        raise ValueError(f"[未知增强操作] op={op}")

def _normalize_probs(op_probs: Dict[str, float], ops: Iterable[str]) -> Dict[str, float]:
    """对给定 ops 子集进行归一化；未提供的键按0处理。"""
    ops = list(ops)
    clean = {op: float(max(0.0, op_probs.get(op, 0.0))) for op in ops}
    s = sum(clean.values())
    if s <= 0:
        raise ValueError(f"[op_probs非法] 指定的操作 {ops} 权重总和需大于0")
    return {op: (clean[op] / s) for op in ops}

def _choose_ops(
    op_probs: Dict[str, float],
    k: int,
    rng: random.Random
) -> List[str]:
    """
    从给定概率的操作集合里“无放回”抽取 k 个不同操作。
    """
    ops = list(op_probs.keys())
    if not ops:
        return []
    k = max(1, min(int(k), len(ops)))
    weights = [op_probs[o] for o in ops]
    chosen: List[str] = []
    tmp_ops = ops[:]
    tmp_weights = weights[:]
    for _ in range(k):
        total = sum(tmp_weights)
        if total <= 0:
            break
        # 手写带权采样
        r = rng.random() * total
        acc = 0.0
        idx = 0
        for i, w in enumerate(tmp_weights):
            acc += w
            if r <= acc:
                idx = i
                break
        chosen.append(tmp_ops[idx])
        tmp_weights[idx] = 0.0
    return chosen

# ==============================
# 核心增强入口
# ==============================

def build_enhanced_cls(
    annotations_path: str,
    enhanced_save_dir: str,
    normal_key: str = "正常眼底",
    count_strategy: str = "random",                  # "random" 或 "fixed"
    per_image_k: int = 2,                            # fixed 时每图增强张数（1~len(allowed_ops)）
    op_probs: Optional[Dict[str, float]] = None,     # 概率字典，对齐 allowed_ops
    seed: int = 42,
    force_rebuild: bool = False,
    # --- 新增：类条件“不增强”概率 ---
    normal_no_aug_prob: float = 0.7,                 # 正常样本不增强概率（默认 70% 不增强）
    abnormal_no_aug_prob: float = 0.0,               # 异常样本不增强概率（默认 0%，即必增强）
    # --- 新增：可用操作子集与随机次数集合 ---
    allowed_ops: Optional[List[str]] = None,         # 可选子集，默认全部["h","v","vh","rot90","rot270","d","ad"]
    random_k_choices: Tuple[int, ...] = (1, 2, 3)    # random 策略下的 k 候选集合
) -> Dict[str, dict]:
    """
    构建增强分类数据集。

    原始标注格式:
        {
          "xxx.png": {
              "image_path": ".../xxx.png",
              "label": ["轻度糖尿病视网膜病变"] 或 ["正常眼底"] 或 None,
              "quality": "良好" 或 None
          },
          ...
        }

    返回:
        {
          "annotations": { 新文件名: {...}, ... },
          "stats": {
              "total_input": int,
              "normal": int,
              "abnormal": int,
              "skipped_no_aug_normal": int,   # 因概率跳过增强的正常样本数
              "skipped_no_aug_abnormal": int, # 因概率跳过增强的异常样本数（通常为0）
              "augmented_images": int,
              "by_op": {"h": n, "v": n, "vh": n, "rot90": n, "rot270": n, "d": n, "ad": n}
          }
        }
    """
    # 准备随机性
    rng = random.Random(seed)

    # 可用操作集（D4的7个非恒等变换）
    ALL_OPS = ["h", "v", "vh", "rot90", "rot270", "d", "ad"]
    if allowed_ops is None:
        allowed_ops = ALL_OPS[:]
    else:
        illegal = [op for op in allowed_ops if op not in ALL_OPS]
        if illegal:
            raise ValueError(f"[allowed_ops非法] 不支持的操作: {illegal}；支持: {ALL_OPS}")

    # 概率校验与归一化
    if op_probs is None:
        # 合理默认：偏向水平/垂直/±90，适量对角线与180°
        op_probs = {"h":0.25,"v":0.25,"vh":0.1,"rot90":0.2,"rot270":0.1,"d":0.05,"ad":0.05}
    op_probs = _normalize_probs(op_probs, allowed_ops)

    # 目录与现有结果处理
    img_dir = os.path.join(enhanced_save_dir, "images")
    ann_path = os.path.join(enhanced_save_dir, "enhanced_annotations.json")
    meta_path = os.path.join(enhanced_save_dir, "meta.json")
    if force_rebuild and os.path.isdir(enhanced_save_dir):
        shutil.rmtree(enhanced_save_dir)
    _safe_makedirs(img_dir)

    if (not force_rebuild) and os.path.exists(ann_path) and os.path.exists(meta_path):
        return {
            "annotations": _safe_load_json(ann_path),
            "stats": _safe_load_json(meta_path).get("stats", {})
        }

    # 读取原始标注
    annotations = _safe_load_json(annotations_path)
    if not isinstance(annotations, dict) or len(annotations) == 0:
        raise ValueError(f"[原始标注为空或格式异常] annotations_path={annotations_path}")

    total_input = len(annotations)
    normal_cnt = 0
    abnormal_cnt = 0
    skipped_no_aug_normal = 0
    skipped_no_aug_abnormal = 0

    enhanced_annotations: Dict[str, dict] = {}
    op_counter = {op: 0 for op in ALL_OPS}

    # 遍历样本
    for image_name, item in annotations.items():
        try:
            img_path = item.get("image_path", None)
            if not img_path or (not os.path.exists(img_path)):
                raise FileNotFoundError(f"原图不存在: {img_path}")

            label = item.get("label", None)  # List[str] 或 None
            quality = item.get("quality", None)
            is_normal = (isinstance(label, list) and len(label) == 1 and label[0] == normal_key)

            if is_normal:
                normal_cnt += 1
                no_aug_prob = float(normal_no_aug_prob)
            else:
                abnormal_cnt += 1
                no_aug_prob = float(abnormal_no_aug_prob)

            # 类条件“不增强”概率：命中则跳过该样本
            if rng.random() < no_aug_prob:
                if is_normal:
                    skipped_no_aug_normal += 1
                else:
                    skipped_no_aug_abnormal += 1
                continue

            # 确定本图增强数量 k
            if count_strategy == "fixed":
                k = max(1, min(int(per_image_k), len(allowed_ops)))
            elif count_strategy == "random":
                if (not isinstance(random_k_choices, tuple)) or (len(random_k_choices) == 0):
                    raise ValueError("[random_k_choices非法] 需为非空元组，如 (1,2,3)")
                # clip 到 1..len(allowed_ops)
                candidates = tuple(sorted(set(max(1, min(int(x), len(allowed_ops))) for x in random_k_choices)))
                k = rng.choice(candidates)
            else:
                raise ValueError(f"[count_strategy非法] 仅支持 'random' 或 'fixed'，当前: {count_strategy}")

            chosen_ops = _choose_ops(op_probs, k, rng)

            with Image.open(img_path) as img:
                img = img.convert("RGB")
                for op in chosen_ops:
                    out_img = _apply_op(img, op)
                    new_name = _build_new_name(image_name, op)
                    abs_save_path = os.path.join(img_dir, new_name)
                    out_img.save(abs_save_path, format="PNG")

                    enhanced_annotations[new_name] = {
                        "orig_image_name": image_name,
                        "orig_dataset": os.path.basename(os.path.dirname(os.path.dirname(img_path))) if os.path.dirname(img_path) else "",
                        "orig_image_path": img_path,
                        "aug_ops": op,
                        "label": label,
                        "quality": quality,
                        "aug_image_path": abs_save_path,
                        # 可选校验： "md5": _md5_of_file(abs_save_path)
                    }
                    op_counter[op] += 1

        except Exception as e:
            raise RuntimeError(f"[增强失败] image_name={image_name}, error={repr(e)}")

    stats = {
        "total_input": total_input,
        "normal": normal_cnt,
        "abnormal": abnormal_cnt,
        "skipped_no_aug_normal": skipped_no_aug_normal,
        "skipped_no_aug_abnormal": skipped_no_aug_abnormal,
        "augmented_images": sum(op_counter.values()),
        "by_op": op_counter
    }

    # 输出
    _safe_dump_json(enhanced_annotations, ann_path)
    _safe_dump_json(
        {
            "source_annotations": os.path.abspath(annotations_path),
            "stats": stats,
            "config": {
                "normal_key": normal_key,
                "count_strategy": count_strategy,
                "per_image_k": per_image_k,
                "op_probs": {op: op_probs.get(op, 0.0) for op in allowed_ops},
                "seed": seed,
                "normal_no_aug_prob": normal_no_aug_prob,
                "abnormal_no_aug_prob": abnormal_no_aug_prob,
                "allowed_ops": allowed_ops,
                "random_k_choices": random_k_choices
            }
        },
        meta_path
    )

    return {"annotations": enhanced_annotations, "stats": stats}
